
import os
from loguru import logger
from langchain_core.messages import HumanMessage, AIMessage, RemoveMessage
from langchain_core.messages.utils import count_tokens_approximately

from core.graph import workflow


def run_demo_with_monitoring():
    """基于LangMem上下文消息压缩的示例."""
    logger.info("🚀 LangGraph 智能体与 LangMem 摘要节点演示")
    logger.info("=" * 60)
    logger.info("📌 当超过1000个令牌时，上下文压缩会自动进行。")
    logger.info("=" * 60)

    # Enable LangSmith tracing if API key is set
    if os.getenv("LANGCHAIN_API_KEY"):
        os.environ["LANGCHAIN_TRACING_V2"] = "true"
        os.environ["LANGCHAIN_PROJECT"] = "langgraph-langmem-demo"
        logger.info("✅ LangSmith 追踪已启用")
        logger.info(f"📊 追踪查看：https://smith.langchain.com")
    else:
        logger.info("⚠️  LangSmith 追踪不可用 (须在.env中设置 LANGCHAIN_API_KEY)")
    logger.info("=" * 60)

    agent = workflow
    config = {"configurable": {"thread_id": "langmem-demo"}}

    queries = [
        "现在几点了？",
        "纽约的天气怎么样？",
        "计算2的10次方",
        "告诉我你所知道的关于Python编程的一切",
        "2500的15%是多少？",
        "搜索关于人工智能的详细信息",
        "伦敦和东京的天气如何？",
        "计算7的阶乘",
        "详细告诉我关于气候变化的情况",
        "前20个质数的和是多少？",
        "比较巴黎和悉尼的天气",
        "529的平方根是多少？"
    ]

    logger.info("📊 开始对话循环 ...")

    for i, query in enumerate(queries, 1):
        logger.info(f"{'=' * 60}")
        logger.info(f"💬 Query {i}: {query}")
        logger.info(f"{'=' * 60}")

        messages_before = agent.get_state(config).values.get("messages", [])
        countable_messages_before = [msg for msg in messages_before if not isinstance(msg, RemoveMessage)]
        tokens_before = count_tokens_approximately(countable_messages_before)

        result = agent.invoke(
            {"messages": [HumanMessage(content=query)]},
            config,
            debug=False
        )

        if result["messages"]:
            last_msg = result["messages"][-1]
            if isinstance(last_msg, AIMessage):
                logger.info(f"🤖 助手：{last_msg.content}")

        messages_after = agent.get_state(config).values.get("messages", [])
        countable_messages_after = [msg for msg in messages_after if not isinstance(msg, RemoveMessage)]
        tokens_after = count_tokens_approximately(countable_messages_after)

        if tokens_before > tokens_after and i > 1:
            logger.info(f"🔄 检测到压缩!")
            logger.info(f"📉 Tokens 从 ~{tokens_before} 减少至 ~{tokens_after}")
            logger.info(f"💾 节约 ~{tokens_before - tokens_after} tokens")

        state = agent.get_state(config).values
        if state.get("context", {}).get("running_summary"):
            summary = state["context"]["running_summary"]
            if hasattr(summary, 'content') and summary.content:
                logger.info(f"📝 摘要生成：{len(summary.content)} chars")

        logger.info(f"📏 当前上下文大小： ~{tokens_after} tokens")
        logger.info(f"📨 历史所有消息大小： {len(countable_messages_after)}")

        if tokens_after > 800:
            logger.info(f"⚠️  接近压缩阈值（1000 个令牌）")


if __name__ == "__main__":
    run_demo_with_monitoring()